{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4d871c04-5cff-468b-ae18-11a450364e91",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-11T16:26:20.891785Z",
     "iopub.status.busy": "2022-06-11T16:26:20.891235Z",
     "iopub.status.idle": "2022-06-11T16:26:22.076373Z",
     "shell.execute_reply": "2022-06-11T16:26:22.075693Z",
     "shell.execute_reply.started": "2022-06-11T16:26:20.891712Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import glob\n",
    "import cv2 as cv\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision.utils import save_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "20762392-5d4a-44e0-8365-cb29def9d6cc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:15:16.823029Z",
     "iopub.status.busy": "2022-06-10T16:15:16.822515Z",
     "iopub.status.idle": "2022-06-10T16:15:25.400302Z",
     "shell.execute_reply": "2022-06-10T16:15:25.399561Z",
     "shell.execute_reply.started": "2022-06-10T16:15:16.822992Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = glob.glob('../../data/horse2zebra/train/trainA/*')\n",
    "b = glob.glob('../../data/horse2zebra/train/trainB/*')\n",
    "\n",
    "data1 = np.empty((len(a), 256, 256, 3))\n",
    "for i, j in enumerate(a):\n",
    "    data1[i] = cv.imread(j)\n",
    "data2 = np.empty((len(b), 256, 256, 3))\n",
    "for i, j in enumerate(b):\n",
    "    data2[i] = cv.imread(j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c3b3d82f-7861-4bb7-9763-d231e8795d06",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:15:25.411162Z",
     "iopub.status.busy": "2022-06-10T16:15:25.410662Z",
     "iopub.status.idle": "2022-06-10T16:15:27.225282Z",
     "shell.execute_reply": "2022-06-10T16:15:27.224558Z",
     "shell.execute_reply.started": "2022-06-10T16:15:25.411136Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data1 = (data1[..., [1,2,0]] - 127.5) / 127.5\n",
    "data2 = (data2[..., [1,2,0]] - 127.5) / 127.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a6e0b14-7334-4436-b818-94b2d8b5d352",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:15:35.617425Z",
     "iopub.status.busy": "2022-06-10T16:15:35.616631Z",
     "iopub.status.idle": "2022-06-10T16:15:35.925330Z",
     "shell.execute_reply": "2022-06-10T16:15:35.924596Z",
     "shell.execute_reply.started": "2022-06-10T16:15:35.617381Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataloader1 = torch.utils.data.DataLoader(\n",
    "    torch.tensor(data1, device='cuda').permute(0, 3, 1, 2),\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True)\n",
    "dataloader2 = torch.utils.data.DataLoader(\n",
    "    Tensor(data1).permute(0, 3, 1, 2),\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6f549703-76e4-4337-b7d5-c7c4dc22ef12",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:15:39.473774Z",
     "iopub.status.busy": "2022-06-10T16:15:39.473204Z",
     "iopub.status.idle": "2022-06-10T16:15:39.477213Z",
     "shell.execute_reply": "2022-06-10T16:15:39.476572Z",
     "shell.execute_reply.started": "2022-06-10T16:15:39.473738Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Parse:\n",
    "    epochs = 200\n",
    "    batch_size = 1\n",
    "    lr = 0.0001\n",
    "    size = 256\n",
    "    residual_times = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "26ead90c-a235-4042-a140-5b26421d1293",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:15:41.595308Z",
     "iopub.status.busy": "2022-06-10T16:15:41.594628Z",
     "iopub.status.idle": "2022-06-10T16:15:41.602003Z",
     "shell.execute_reply": "2022-06-10T16:15:41.601270Z",
     "shell.execute_reply.started": "2022-06-10T16:15:41.595270Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.block = nn.Sequential(\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(256, 256, 3),\n",
    "            nn.InstanceNorm2d(256, affine=True),\n",
    "            nn.ReLU(inplace=True), # 加上replace原地操作，那失去了原本的信息，还怎么计算梯度呢\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(256, 256, 3),\n",
    "            nn.InstanceNorm2d(256, affine=True)) # 为什么不再激活\n",
    "\n",
    "    def forward(self, data):\n",
    "        return data + self.block(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "aec2dca5-c9b4-4be7-a984-90c1c06c7e2a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-11T17:05:17.542608Z",
     "iopub.status.busy": "2022-06-11T17:05:17.541785Z",
     "iopub.status.idle": "2022-06-11T17:05:17.551892Z",
     "shell.execute_reply": "2022-06-11T17:05:17.551186Z",
     "shell.execute_reply.started": "2022-06-11T17:05:17.542548Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # 用大视野来转换成特征\n",
    "        model = [\n",
    "            nn.ReflectionPad2d(3),\n",
    "            nn.Conv2d(3, 64, 7),  # nx64x256x256\n",
    "            nn.InstanceNorm2d(64, affine=True),\n",
    "            nn.ReLU(inplace=True)\n",
    "        ]\n",
    "\n",
    "        # 提取特征\n",
    "        model += [\n",
    "            nn.Conv2d(64, 128, 3, stride=2, padding=1),  # nx128x128x128\n",
    "            nn.InstanceNorm2d(128, affine=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(128, 256, 3, stride=2, padding=1),  # nx256x64x64\n",
    "            nn.InstanceNorm2d(256, affine=True),\n",
    "            nn.ReLU(inplace=True)\n",
    "        ]\n",
    "\n",
    "        # 添加残差板块\n",
    "        for _ in range(Parse.residual_times):\n",
    "            model.append(ResidualBlock())\n",
    "\n",
    "        # 下采样\n",
    "        model += [\n",
    "            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),  # nx128x128x128\n",
    "            nn.InstanceNorm2d(128, affine=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),  # nx64x256x256\n",
    "            nn.InstanceNorm2d(64, affine=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ReflectionPad2d(3),\n",
    "            nn.Conv2d(64, 3, 7),\n",
    "            nn.Tanh()\n",
    "        ]\n",
    "\n",
    "        self.model = nn.Sequential(*model)\n",
    "\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=Parse.lr)\n",
    "\n",
    "    def forward(self, data):\n",
    "        return self.model(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "39dc5856-e261-4e69-8624-c699b72de19d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-11T17:13:03.657799Z",
     "iopub.status.busy": "2022-06-11T17:13:03.657096Z",
     "iopub.status.idle": "2022-06-11T17:13:03.664867Z",
     "shell.execute_reply": "2022-06-11T17:13:03.664185Z",
     "shell.execute_reply.started": "2022-06-11T17:13:03.657769Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        model = [\n",
    "            nn.Conv2d(3, 64, 4, stride=2, padding=1),  # 128\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(64, 128, 4, stride=2, padding=1),\n",
    "            nn.InstanceNorm2d(128, affine=True),  # 256->128->64->32->31\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(128, 256, 4, stride=2, padding=1),\n",
    "            nn.InstanceNorm2d(256, affine=True),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(256, 512, 4, padding=1),  # 这里不stride\n",
    "            nn.InstanceNorm2d(512, affine=True),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(512, 1, 4, padding=1)]  # 开始提取特征 31->30\n",
    "\n",
    "        self.model = nn.Sequential(*model)\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=Parse.lr)\n",
    "        \n",
    "    def forward(self, data):\n",
    "        return F.avg_pool2d(self.model(data), 30).squeeze()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ab8b7e14-bdc5-4438-b007-122cfca2df6a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:16:42.129075Z",
     "iopub.status.busy": "2022-06-10T16:16:42.128516Z",
     "iopub.status.idle": "2022-06-10T16:16:42.352741Z",
     "shell.execute_reply": "2022-06-10T16:16:42.352075Z",
     "shell.execute_reply.started": "2022-06-10T16:16:42.129050Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "netD1 = Discriminator()\n",
    "netD1 = Discriminator()\n",
    "netG1 = Generator()\n",
    "netG2 = Generator()\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    netD1.cuda()\n",
    "    netD2.cuda()\n",
    "    netG1.cuda()\n",
    "    netG2.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b272dc9-6448-4e1a-9446-4a5e292ef688",
   "metadata": {},
   "source": [
    "# 三种损失\n",
    "1. 对抗损失\n",
    "2. 循环损失\n",
    "3. 自我损失"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c91d56a0-72a8-41c8-85fd-9af24acacef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_penalty(D, real_samples, fake_samples):\n",
    "    \n",
    "        alpha1 = torch.rand(real_samples.size(0), 1, 1, 1, requires_grad=True, device='cuda')\n",
    "        interpolates = alpha1 * real_samples + ((1 - alpha1) * fake_samples)\n",
    "        d_interpolates = D(interpolates)\n",
    "        fake = torch.ones(real_samples.shape[0], device='cuda')\n",
    "        gradients = autograd.grad(\n",
    "            outputs=d_interpolates,\n",
    "            inputs=interpolates,\n",
    "            grad_outputs=fake,\n",
    "            create_graph=True,\n",
    "            retain_graph=True,\n",
    "            only_inputs=True,\n",
    "        )[0]\n",
    "\n",
    "        gradients = gradients.reshape(gradients.size(0), -1)\n",
    "        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n",
    "        \n",
    "        return gradient_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a6d92d22-8c2c-4ddd-a363-43f1740868b8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-10T16:16:50.166904Z",
     "iopub.status.busy": "2022-06-10T16:16:50.166111Z",
     "iopub.status.idle": "2022-06-10T16:17:03.295561Z",
     "shell.execute_reply": "2022-06-10T16:17:03.294905Z",
     "shell.execute_reply.started": "2022-06-10T16:16:50.166863Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sample_image(i, generator, status, z=None):\n",
    "    if not z:\n",
    "        z = torch.randn(5 ** 2, opt.latent_dim, 1, 1, device=next(generator.parameters()).device)\n",
    "    gen_imgs = generator(z, status, 1)\n",
    "    gen_imgs = F.interpolate(gen_imgs, scale_factor=128*5/gen_imgs.shape[2], mode=\"nearest\")\n",
    "    save_image(gen_imgs.data, f\"images/{i}.png\", nrow=n_row, normalize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "72adac15-13ed-4bfc-9a5b-24aa110edcef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-11T17:56:42.594632Z",
     "iopub.status.busy": "2022-06-11T17:56:42.594375Z",
     "iopub.status.idle": "2022-06-11T17:56:42.599733Z",
     "shell.execute_reply": "2022-06-11T17:56:42.599097Z",
     "shell.execute_reply.started": "2022-06-11T17:56:42.594608Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(torch.randn(5,6).data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "63e3579e-2e9a-4ef2-91b8-d704f2960525",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-11T01:50:34.492515Z",
     "iopub.status.busy": "2022-06-11T01:50:34.491858Z",
     "iopub.status.idle": "2022-06-11T01:50:40.923459Z",
     "shell.execute_reply": "2022-06-11T01:50:40.922799Z",
     "shell.execute_reply.started": "2022-06-11T01:50:34.492479Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mkpic(3, '马')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a4ca88-b678-404d-93db-bed2ded64c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# loss_GAN\n",
    "loss_pix = nn.L1Loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa1f7e2-6331-4604-b40e-54864da4a990",
   "metadata": {
    "collapsed": true,
    "execution": {
     "iopub.execute_input": "2022-06-10T16:17:23.877551Z",
     "iopub.status.busy": "2022-06-10T16:17:23.877007Z"
    },
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch：0\n",
      "0\n",
      "10\n",
      "20\n",
      "30\n",
      "40\n",
      "50\n",
      "60\n",
      "70\n",
      "80\n",
      "90\n",
      "100\n",
      "110\n",
      "120\n",
      "130\n",
      "140\n",
      "150\n",
      "160\n",
      "170\n",
      "180\n",
      "190\n",
      "200\n",
      "210\n",
      "220\n",
      "230\n",
      "240\n",
      "250\n",
      "260\n",
      "270\n",
      "280\n",
      "290\n",
      "300\n",
      "310\n",
      "320\n",
      "330\n",
      "340\n",
      "350\n",
      "360\n",
      "370\n",
      "380\n",
      "390\n",
      "400\n",
      "410\n",
      "420\n",
      "430\n",
      "440\n",
      "450\n",
      "460\n",
      "470\n",
      "480\n",
      "490\n",
      "500\n",
      "510\n",
      "520\n",
      "530\n",
      "540\n",
      "550\n",
      "560\n",
      "570\n",
      "580\n",
      "590\n",
      "600\n",
      "610\n",
      "620\n",
      "630\n",
      "640\n",
      "650\n",
      "660\n",
      "670\n",
      "680\n",
      "690\n",
      "700\n",
      "710\n",
      "720\n",
      "730\n",
      "740\n",
      "750\n",
      "760\n",
      "770\n",
      "780\n",
      "790\n",
      "800\n",
      "810\n",
      "820\n",
      "830\n",
      "840\n",
      "850\n",
      "860\n",
      "870\n",
      "880\n",
      "890\n",
      "900\n",
      "910\n",
      "920\n",
      "930\n",
      "940\n",
      "950\n",
      "960\n",
      "970\n",
      "980\n",
      "990\n",
      "1000\n",
      "1010\n",
      "1020\n",
      "1030\n",
      "1040\n",
      "1050\n",
      "1060\n",
      "epoch：1\n",
      "0\n",
      "10\n",
      "20\n",
      "30\n",
      "40\n",
      "50\n",
      "60\n",
      "70\n",
      "80\n",
      "90\n",
      "100\n",
      "110\n",
      "120\n",
      "130\n",
      "140\n",
      "150\n",
      "160\n",
      "170\n",
      "180\n",
      "190\n",
      "200\n",
      "210\n",
      "220\n",
      "230\n",
      "240\n",
      "250\n",
      "260\n",
      "270\n",
      "280\n",
      "290\n",
      "300\n",
      "310\n",
      "320\n",
      "330\n",
      "340\n",
      "350\n",
      "360\n",
      "370\n",
      "380\n",
      "390\n",
      "400\n",
      "410\n",
      "420\n",
      "430\n",
      "440\n",
      "450\n",
      "460\n",
      "470\n",
      "480\n",
      "490\n",
      "500\n",
      "510\n",
      "520\n",
      "530\n",
      "540\n",
      "550\n",
      "560\n",
      "570\n",
      "580\n",
      "590\n",
      "600\n",
      "610\n",
      "620\n",
      "630\n",
      "640\n",
      "650\n",
      "660\n",
      "670\n",
      "680\n",
      "690\n",
      "700\n",
      "710\n",
      "720\n",
      "730\n",
      "740\n",
      "750\n",
      "760\n",
      "770\n",
      "780\n",
      "790\n",
      "800\n",
      "810\n",
      "820\n",
      "830\n",
      "840\n",
      "850\n",
      "860\n",
      "870\n",
      "880\n",
      "890\n",
      "900\n",
      "910\n",
      "920\n",
      "930\n",
      "940\n",
      "950\n",
      "960\n",
      "970\n",
      "980\n",
      "990\n",
      "1000\n",
      "1010\n",
      "1020\n",
      "1030\n",
      "1040\n",
      "1050\n",
      "1060\n",
      "epoch：2\n",
      "0\n",
      "10\n",
      "20\n",
      "30\n",
      "40\n",
      "50\n",
      "60\n",
      "70\n",
      "80\n",
      "90\n",
      "100\n",
      "110\n",
      "120\n",
      "130\n",
      "140\n",
      "150\n",
      "160\n",
      "170\n",
      "180\n",
      "190\n",
      "200\n",
      "210\n",
      "220\n",
      "230\n",
      "240\n",
      "250\n",
      "260\n",
      "270\n",
      "280\n",
      "290\n",
      "300\n",
      "310\n",
      "320\n",
      "330\n",
      "340\n",
      "350\n",
      "360\n",
      "370\n",
      "380\n",
      "390\n",
      "400\n",
      "410\n",
      "420\n",
      "430\n",
      "440\n",
      "450\n",
      "460\n",
      "470\n",
      "480\n",
      "490\n",
      "500\n",
      "510\n",
      "520\n",
      "530\n",
      "540\n",
      "550\n",
      "560\n",
      "570\n",
      "580\n",
      "590\n",
      "600\n",
      "610\n",
      "620\n",
      "630\n",
      "640\n",
      "650\n",
      "660\n",
      "670\n",
      "680\n",
      "690\n",
      "700\n",
      "710\n",
      "720\n",
      "730\n",
      "740\n",
      "750\n",
      "760\n",
      "770\n",
      "780\n",
      "790\n",
      "800\n",
      "810\n",
      "820\n",
      "830\n",
      "840\n",
      "850\n",
      "860\n",
      "870\n",
      "880\n",
      "890\n",
      "900\n",
      "910\n",
      "920\n",
      "930\n",
      "940\n",
      "950\n",
      "960\n",
      "970\n",
      "980\n",
      "990\n",
      "1000\n",
      "1010\n",
      "1020\n",
      "1030\n",
      "1040\n",
      "1050\n",
      "1060\n",
      "epoch：3\n",
      "0\n",
      "10\n",
      "20\n",
      "30\n",
      "40\n",
      "50\n",
      "60\n",
      "70\n",
      "80\n",
      "90\n",
      "100\n",
      "110\n",
      "120\n",
      "130\n",
      "140\n",
      "150\n",
      "160\n",
      "170\n",
      "180\n",
      "190\n",
      "200\n",
      "210\n",
      "220\n",
      "230\n",
      "240\n",
      "250\n",
      "260\n",
      "270\n",
      "280\n",
      "290\n",
      "300\n",
      "310\n",
      "320\n",
      "330\n",
      "340\n",
      "350\n",
      "360\n",
      "370\n",
      "380\n",
      "390\n",
      "400\n",
      "410\n",
      "420\n",
      "430\n",
      "440\n",
      "450\n",
      "460\n",
      "470\n",
      "480\n",
      "490\n",
      "500\n",
      "510\n",
      "520\n",
      "530\n",
      "540\n",
      "550\n",
      "560\n",
      "570\n",
      "580\n",
      "590\n",
      "600\n",
      "610\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(Parse.epochs):\n",
    "\n",
    "    print(f'epoch：{ epoch }')\n",
    "\n",
    "    for i, (realA, realB) in enumerate(zip(dataloader1, dataloader2)):\n",
    "        if not i % 10:\n",
    "            print(i)\n",
    "\n",
    "        sameA = netG1(realA)\n",
    "        fakeB = netG1(realA)\n",
    "        recoverA = netG2(fakeB)\n",
    "\n",
    "        sameB = netG2(realB)\n",
    "        fakeA = netG2(realB)\n",
    "        recoverB = netG1(fakeA)\n",
    "        \n",
    "        critic2 = netD2(fakeB)\n",
    "        critic1 = netD1(fakeA)\n",
    "        \n",
    "        critic1_r = netD1(realA)\n",
    "        critic2_r = netD2(realB)\n",
    "        \n",
    "        # 训练D\n",
    "        netD1.optimizer.zero_grad()\n",
    "        netD2.optimizer.zero_grad()\n",
    "        loss1 = - critic1_r + critic1 + gradient_penalty(netD1, realA, fakeA)\n",
    "        loss2 = - critic2_r + critic2 + gradient_penalty(netD2, realB, fakeB)\n",
    "        loss1.backward()\n",
    "        loss2.backward()\n",
    "        netD1.optimizer.step()\n",
    "        netD2.optimizer.step()\n",
    "        \n",
    "        \n",
    "        netG1.optimizer.zero_grad()\n",
    "        netG2.optimizer.zero_grad()\n",
    "        loss1 = - critic1 + loss_pix(sameA, realA) + loss_pix(recoverA, realA)\n",
    "        loss2 = - critic2 + loss_pix(sameB, realB) + loss_pix(recoverB, realB)\n",
    "        loss1.backward()\n",
    "        loss2.backward()\n",
    "        netG1.optimizer.step()\n",
    "        netG2.optimizer.step()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a104ba2-7934-4d83-9729-b326e27e1fd2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96cb46da-ace0-4f76-be96-a1a840e0690b",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.init.normal_(tensor, mean=0, std=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
